package org.xi.dsaa.collection;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;

/**
 * 红黑树 Red Black Tree
 *
 * @param <K> 元素类型
 */
public class RedBlackTreeMap<K extends Comparable<K>, V> {

    private static final boolean RED = false;
    private static final boolean BLACK = true;

    private Entry<K, V> root;

    public Entry<K, V> getRoot() {
        return root;
    }

    public void setRoot(Entry<K, V> root) {
        this.root = root;
    }

    public int size() {
        return 0;
    }

    public boolean isEmpty() {
        return root == null;
    }

    public V put(K key, V value) {
        if (key == null) throw new NullPointerException("");

        if (root == null) root = new Entry<>(key, value);

        Entry<K, V> parent, node = root;
        int cmp;

        do {
            parent = node;
            cmp = node.key.compareTo(key);
            if (cmp == 0) {
                V oldValue = node.value;
                node.setValue(value);
                return oldValue;
            } else if (cmp > 0) {
                node = node.left;
            } else {
                node = node.right;
            }
        } while (node != null);

        node = new Entry<>(key, value, parent);
        if (cmp > 0) {
            parent.left = node;
        } else {
            parent.right = node;
        }
        fixAfterPut(node);
        return null;
    }

    private void fixAfterPut(Entry<K, V> node) {
        node.color = RED;
        while (node != null && node != root && node.parent.color == RED) {
            if (parentOf(node) == leftOf(parentOf(parentOf(node)))) {

                Entry<K, V> parentRight = rightOf(parentOf(parentOf(node)));
                if (colorOf(parentRight) == RED) {
                    setColor(parentOf(node), BLACK);
                    setColor(parentRight, BLACK);
                    setColor(parentOf(parentOf(node)), BLACK);
                    node = parentOf(parentOf(node));
                } else {
                    if (node == rightOf(parentOf(node))) {
                        node = parentOf(node);
                        rotateLeft(node);
                    }
                    setColor(parentOf(node), BLACK);
                    setColor(parentOf(parentOf(node)), RED);
                    rotateRight(parentOf(parentOf(node)));
                }
            } else {
                Entry<K, V> parentLeft = leftOf(parentOf(parentOf(node)));
                if (colorOf(parentLeft) == RED) {
                    setColor(parentOf(node), BLACK);
                    setColor(parentLeft, BLACK);
                    setColor(parentOf(parentOf(node)), RED);
                    node = parentOf(parentOf(node));
                } else {
                    if (node == leftOf(parentOf(node))) {
                        node = parentOf(node);
                        rotateRight(node);
                    }
                    setColor(parentOf(node), BLACK);
                    setColor(parentOf(parentOf(node)), RED);
                    rotateLeft(parentOf(parentOf(node)));
                }
            }
        }
        root.color = BLACK;
    }

    public V remove(Object key) {
        Entry<K, V> p = getEntry((K) key);
        if (p == null)
            return null;

        V oldValue = p.value;
        deleteEntry(p);
        return oldValue;
    }

    private void deleteEntry(Entry<K, V> node) {
        // 左子节点和右子节点都不为空，找后继节点代替当前节点，将要删除的节点指向后继节点的位置
        if (node.left != null && node.right != null) {
            Entry<K, V> successor = successor(node);
            node.key = successor.key;
            node.value = successor.value;
            node = successor;
        }
        // 此时，node为叶子节点，或者只有左叶子节点/右叶子节点的上级节点
        Entry<K, V> replaceNode = node.left == null ? node.right : node.left;
        if (replaceNode != null) {
            // node节点为叶子节点的上级节点
            // 用子节点替换当前节点
            replaceNode.parent = node.parent;
            if (node.parent == null) {
                // node的父节点为null，说明node是根结点，将替换的子节点作为新的根结点
                root = replaceNode;
            } else if (node == node.parent.left) {
                // node是父节点的左子节点，替换的子节点依然是左子节点
                node.parent.left = replaceNode;
            } else {
                // node是父节点的右子节点，替换的子节点依然是右子节点
                node.parent.right = replaceNode;
            }
            // 垃圾回收
            node.parent = node.left = node.right = null;
            if (node.color == BLACK) {
                // 替代节点是红色，只需要变色
                fixAfterRemove(replaceNode);
            }
        } else if (node.parent == null) {
            // node节点是叶子节点，node的父节点为null，说明node是根结点，此时删除的是根结点
            root = null;
        } else {
            // node节点是叶子节点
            if (node.color == BLACK) {
                // 先调整
                fixAfterRemove(node);
            }

            // 再删除
            if (node.parent != null) {
                if (node == node.parent.left) {
                    node.parent.left = null;
                } else if (node == node.parent.right) {
                    node.parent.right = null;
                }
                node.parent = null;
            }
        }
    }

    private void fixAfterRemove(Entry<K, V> node) {
        while (node != root && node.color == BLACK) {
            if (node == leftOf(parentOf(node))) {
                // node是左子节点
                Entry<K, V> sib = rightOf(parentOf(node));
                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(node), RED);
                    rotateLeft(parentOf(node));
                    sib = rightOf(parentOf(node));
                }
                if (colorOf(leftOf(sib)) == BLACK && colorOf(rightOf(sib)) == BLACK) {
                    // 兄弟节点左子树和右子树都为空
                    setColor(sib, RED);
                    node = parentOf(node);
                } else {
                    if (colorOf(rightOf(sib)) == BLACK) {
                        // 兄弟节点右子树为空，左子树不为空
                        setColor(leftOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateRight(sib);
                        sib = rightOf(parentOf(node));
                    }
                    setColor(sib, colorOf(parentOf(node)));
                    setColor(parentOf(node), BLACK);
                    setColor(rightOf(sib), BLACK);
                    rotateLeft(parentOf(node));
                    node = root;
                }
            } else {
                // node是右子节点
                Entry<K, V> sib = leftOf(parentOf(node));
                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(node), RED);
                    rotateRight(parentOf(node));
                    sib = leftOf(parentOf(node));
                }
                if (colorOf(rightOf(sib)) == BLACK && colorOf(leftOf(sib)) == BLACK) {
                    // 兄弟节点左子树和右子树都为空
                    setColor(sib, RED);
                    node = parentOf(node);
                } else {
                    if (colorOf(leftOf(sib)) == BLACK) {
                        // 兄弟节点左子树为空，右子树不为空
                        setColor(rightOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateLeft(sib);
                        sib = leftOf(parentOf(node));
                    }
                    setColor(sib, colorOf(parentOf(node)));
                    setColor(parentOf(node), BLACK);
                    setColor(leftOf(sib), BLACK);
                    rotateRight(parentOf(node));
                    node = root;
                }
            }
        }
        setColor(node, BLACK);
    }

    static <K extends Comparable<K>, V> Entry<K, V> successor(Entry<K, V> t) {
        if (t == null)
            return null;
        else if (t.right != null) {
            Entry<K, V> p = t.right;
            while (p.left != null)
                p = p.left;
            return p;
        } else {
            Entry<K, V> p = t.parent;
            Entry<K, V> ch = t;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    static <K extends Comparable<K>, V> Entry<K, V> predecessor(Entry<K, V> t) {
        if (t == null)
            return null;
        else if (t.left != null) {
            Entry<K, V> p = t.left;
            while (p.right != null)
                p = p.right;
            return p;
        } else {
            Entry<K, V> p = t.parent;
            Entry<K, V> ch = t;
            while (p != null && ch == p.left) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    private static <K extends Comparable<K>, V> boolean colorOf(Entry<K, V> node) {
        return (node == null ? BLACK : node.color);
    }

    private static <K extends Comparable<K>, V> Entry<K, V> parentOf(Entry<K, V> node) {
        return (node == null ? null : node.parent);
    }

    private static <K extends Comparable<K>, V> void setColor(Entry<K, V> node, boolean color) {
        if (node != null) node.color = color;
    }

    private static <K extends Comparable<K>, V> Entry<K, V> leftOf(Entry<K, V> node) {
        return (node == null) ? null : node.left;
    }

    private static <K extends Comparable<K>, V> Entry<K, V> rightOf(Entry<K, V> node) {
        return (node == null) ? null : node.right;
    }

    final Entry<K, V> getMin(Entry<K, V> node) {

        Entry<K, V> parent = null;
        while (node != null) {
            parent = node;
            node = node.left;
        }
        return parent;
    }

    final Entry<K, V> getMax(Entry<K, V> node) {
        Entry<K, V> parent = null;
        while (node != null) {
            parent = node;
            node = node.right;
        }
        return parent;
    }

    final Entry<K, V> getEntry(K key) {
        if (key == null)
            throw new NullPointerException();
        Entry<K, V> cursor = root;
        while (cursor != null) {
            int cmp = key.compareTo(cursor.key);
            if (cmp < 0)
                cursor = cursor.left;
            else if (cmp > 0)
                cursor = cursor.right;
            else
                return cursor;
        }
        return null;
    }

    // 将entry节点的右子节点的左子节点作为entry节点的右子节点，将entry节点作为原右子节点的左子节点
    private void rotateLeft(Entry<K, V> node) {
        if (node == null) return;

        Entry<K, V> parent = node.parent;
        Entry<K, V> right = node.right;

        node.right = right.left;
        if (node.right != null) node.right.parent = node;

        right.left = node;
        node.parent = right;

        right.parent = parent;
        // 父节点为空，说明entry是根结点，将entry的右子节点作为新的根结点，返回
        if (parent == null) {
            root = right;
            return;
        }

        if (parent.right == node) {
            parent.right = right;
        } else {
            parent.left = right;
        }
    }

    private void rotateRight(Entry<K, V> node) {
        if (node == null) return;

        Entry<K, V> parent = node.parent;
        Entry<K, V> left = node.left;

        node.left = left.right;
        if (node.left != null) node.left.parent = node;

        left.right = node;
        node.parent = left;

        left.parent = parent;
        // 父节点为空，说明entry是根结点，将entry的右子节点作为新的根结点，返回
        if (parent == null) {
            root = left;
            return;
        }

        if (parent.right == node) {
            parent.right = left;
        } else {
            parent.left = left;
        }
    }

    @NoArgsConstructor
    @AllArgsConstructor
    @Data
    public static final class Entry<K extends Comparable<K>, V> {

        public Entry(K key, V value) {
            this(key, value, null, null, null, RED);
        }

        public Entry(K key, V value, boolean color) {
            this(key, value, null, null, null, color);
        }

        public Entry(K key, V value, Entry<K, V> parent) {
            this(key, value, null, null, parent, RED);
        }

        K key;
        V value;
        Entry<K, V> left;
        Entry<K, V> right;
        Entry<K, V> parent;
        boolean color = BLACK;
    }
}
